package cn.iocoder.yudao.module.digital.utils;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.stream.Collectors;

public class AIRecognizeUtil {

    private static final String API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation";
    private static final String API_KEY = "sk-deda9717bbb54c0bbee0b6d98550843e";

    /**
     * 文件识别信息
     * @param imageBytes 文件资源
     * @param additionalExtraction 附件额外提取条件信息提示词
     * @return
     * @throws IOException
     */
    public static String recognizeImage(byte[] imageBytes, String additionalExtraction) throws IOException {
        String mimeType = getImageMimeType(imageBytes);
        String base64Image = Base64.getEncoder().encodeToString(imageBytes);
        String dataUri = "data:" + mimeType + ";base64," + base64Image;

        String prompt = "请从下方图片内容中识别并提取关键信息，并以 JSON 格式返回（不要有自然语言说明、不要换行、不要加解释、只返回 JSON 对象）："
                + "请根据图片内容判断是以下哪种类型：论文/专利/软件著作权（软著）/项目/竞赛，并提取对应的字段。"
                + "1. 如果是【论文】，请提取字段：'论文题目','发表/出版时间','发表刊物/论文集','所属院系','ISSN号','版面','CN号'。"
                + "2. 如果是【专利】，请提取字段：'专利名称','所属院系','专利类型（发明专利/实用新型/外观设计）','申请日期','申请号','授权号','申请人','发明人'。"
                + "3. 如果是【软件著作权（软著）】，请提取字段：'著作权名称','著作权编号','所属院系','登记号','出版日期','作者名称'。"
                + "4. 如果是【项目】，请提取字段：'项目名称','项目类型（教改/教研/科研）','主办方','级别（校级/市级/省级/国家级）','项目成员','立项时间','结项时间','状态（未结项/已结项）'。"
                + "5. 如果是【竞赛】，请提取字段：'竞赛名称','作品名称','参赛学生','比赛级别（查询资料判断此比赛属于什么级别，如：A+、A、B+、B等）','奖项（如几等奖）','指导老师','获奖时间'。"
                + "其他通用字段若有出现也请提取（注意：不要省略这里的任何字段）：'学校','证书编号','发证单位','发证日期','认证机构','主办单位','参赛队伍','参赛队员','比赛名称'。"
                + "附加字段也需要提取（如果与上面的字段有所冲突，提取规则以附加字段为准）：" + additionalExtraction
                + "上面所有字段中如果某个字段在文档中没有出现，请返回'字段名':''。"
                + "注意：如果字段后面有括号的是必须从中选定一个作为字段的值，查找图片中有无写明，若无自行分析内容判断括号中哪个最符合，选定其为字段的值，当然最后所有返回字段必须去掉括号和括号里的内容。";

        String requestBody = "{\n" +
                "  \"model\": \"qwen-vl-max\",\n" +
                "  \"input\": {\n" +
                "    \"messages\": [\n" +
                "      {\n" +
                "        \"role\": \"user\",\n" +
                "        \"content\": [\n" +
                "          {\"text\": \"" + prompt.replace("\"", "\\\"") + "\"},\n" +
                "          {\"image\": \"" + dataUri + "\"}\n" +
                "        ]\n" +
                "      }\n" +
                "    ]\n" +
                "  },\n" +
                "  \"parameters\": {\n" +
                "    \"ocr\": {\n" +
                "      \"enable\": true\n" +
                "    }\n" +
                "  }\n" +
                "}";

        URL url = new URL(API_URL);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Authorization", "Bearer " + API_KEY);
        conn.setRequestProperty("Content-Type", "application/json");
        conn.setDoOutput(true);
        conn.setConnectTimeout(10000);
        conn.setReadTimeout(20000);

        try (OutputStream os = conn.getOutputStream()) {
            os.write(requestBody.getBytes(StandardCharsets.UTF_8));
        }

        InputStream is = (conn.getResponseCode() >= 200 && conn.getResponseCode() < 300)
                ? conn.getInputStream()
                : conn.getErrorStream();

        try (BufferedReader reader = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) {
            return reader.lines().collect(Collectors.joining("\n"));
        }
    }

    public static byte[] convertPdfToImageBytes(InputStream pdfInputStream) throws IOException {
        PDDocument document = PDDocument.load(pdfInputStream);
        PDFRenderer renderer = new PDFRenderer(document);

        // 仅识别第一页，你也可以改为合并全部页
        BufferedImage image = renderer.renderImageWithDPI(0, 300); // 高分辨率
        document.close();

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(image, "jpg", baos); // 输出为 JPEG
        return baos.toByteArray();
    }

    private static String getImageMimeType(byte[] imageData) {
        if (imageData.length >= 4) {
            if ((imageData[0] & 0xFF) == 0xFF && (imageData[1] & 0xFF) == 0xD8) {
                return "image/jpeg";
            }
            if ((imageData[0] & 0xFF) == 0x89 && (imageData[1] & 0xFF) == 0x50) {
                return "image/png";
            }
        }
        return "image/jpeg"; // 默认
    }

    /**
     * 提取 doc 文本内容
     */
    public static String extractTextFromDoc(InputStream in) throws IOException {
        try (HWPFDocument doc = new HWPFDocument(in)) {
            WordExtractor extractor = new WordExtractor(doc);
            return extractor.getText();
        }
    }

    /**
     * 提取 docx 文本内容
     */
    public static String extractTextFromDocx(InputStream in) throws IOException {
        try (XWPFDocument docx = new XWPFDocument(in)) {
            XWPFWordExtractor extractor = new XWPFWordExtractor(docx);
            return extractor.getText();
        }
    }

    /**
     * 提取 txt 文本内容
     */
    public static String extractTextFromTxt(InputStream in) throws IOException {
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))) {
            return reader.lines().collect(Collectors.joining("\n"));
        }
    }

    /**
     * 识别文本内容的关键信息（使用org.json构造JSON，避免转义错误）
     */
    public static String recognizeText(String text) throws IOException {
        String apiUrl = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions";
        String apiKey = API_KEY;
        String prompt = "请从下方文本中识别并提取关键信息，并以 JSON 格式返回（不要有自然语言说明、不要换行、不要加解释、只返回 JSON 对象）：\n\n"
                + "请根据文档内容判断是以下哪种类型：论文 / 专利 / 软件著作权（软著） / 项目 / 竞赛，并提取对应的字段：\n\n"

                + "1. 如果是【论文】，请提取字段：\n"
                + "\"论文题目\", \"发表/出版时间\", \"发表刊物/论文集\", \"所属院系\", \"ISSN号\", \"版面\", \"CN号\"\n\n"

                + "2. 如果是【专利】，请提取字段：\n"
                + "\"专利名称\", \"所属院系\", \"专利类型（发明专利/实用新型/外观设计）\", \"申请日期\", \"申请号\", \"授权号\", \"申请人\", \"发明人\"\n\n"

                + "3. 如果是【软件著作权（软著）】，请提取字段：\n"
                + "\"著作权名称\", \"著作权编号\", \"所属院系\", \"登记号\", \"出版日期\", \"作者名称\"\n\n"

                + "4. 如果是【项目】，请提取字段：\n"
                + "\"项目名称\", \"项目类型（教改/教研/科研）\", \"主办方\", \"级别（校级/市级/省级/国家级）\", \"项目成员\", \"立项时间\", \"结项时间\", \"状态（未结项/已结项）\"\n\n"

                + "5. 如果是【竞赛】，请提取字段：\n"
                + "\"竞赛名称\", \"作品名称\", \"参赛学生\", \"比赛级别（查询资料判断此比赛属于什么级别，如：A+、A、B+、B等）\", \"奖项（如几等奖）\", \"指导老师\", \"获奖时间\"\n\n"

                + "其他通用字段若有出现也请提取：\n"
                + "\"学校\", \"证书编号\", \"发证单位\", \"发证日期\", \"认证机构\", \"主办单位\", \"参赛队伍\", \"参赛队员\", \"比赛名称\"\n\n"

                + "上面所有字段中如果某个字段在文档中没有出现，请返回空字符串。\n"
                + "注意：如果字段后面有括号的是必须从中选定一个作为字段的值，查找图片中有无写明，若无自行分析内容判断括号中哪个最符合，选定其为字段的值，当然最后所有返回字段必须去掉括号和括号里的内容。"
                + text;

        org.json.JSONObject requestBody = new org.json.JSONObject();
        requestBody.put("model", "qwen-plus");
        org.json.JSONArray messages = new org.json.JSONArray();
        org.json.JSONObject systemMsg = new org.json.JSONObject();
        systemMsg.put("role", "system");
        systemMsg.put("content", "You are a helpful assistant.");
        messages.put(systemMsg);
        org.json.JSONObject userMsg = new org.json.JSONObject();
        userMsg.put("role", "user");
        userMsg.put("content", prompt);
        messages.put(userMsg);
        requestBody.put("messages", messages);

        URL url = new URL(apiUrl);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Authorization", "Bearer " + apiKey);
        conn.setRequestProperty("Content-Type", "application/json");
        conn.setDoOutput(true);
        conn.setConnectTimeout(10000);
        conn.setReadTimeout(20000);

        try (OutputStream os = conn.getOutputStream()) {
            os.write(requestBody.toString().getBytes(StandardCharsets.UTF_8));
        }

        InputStream is = (conn.getResponseCode() >= 200 && conn.getResponseCode() < 300)
                ? conn.getInputStream()
                : conn.getErrorStream();

        try (BufferedReader reader = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) {
            return reader.lines().collect(Collectors.joining("\n"));
        }
    }
}